import json
from multiprocessing import Process, cpu_count

from kafka import KafkaConsumer
import redis
from loguru import logger

import config

COMMIT_BUFFER_SIZE = 1
# 重复消费控制
FILTER_SWITCH = False

class Consumer():

    def __init__(self,
                 topic,
                 bootstrap_servers=config.SERVERS,
                 group_id=config.GROUP_ID):
        self.kafka_consumer = KafkaConsumer(
            topic,
            bootstrap_servers=bootstrap_servers,
            group_id=group_id,
            enable_auto_commit=False,
            auto_offset_reset='earliest',
            value_deserializer=lambda msg: json.loads(msg.decode()))
        self.cache = redis.Redis(
            host='localhost',  # ip地址
            port=6379,  # 端口号，默认为6379
            db=0,
        ) if FILTER_SWITCH else None
        self.topic = topic

    def consume(self, handler):
        count = 0
        for msg in self.kafka_consumer:
            self._consume(msg, handler)
            count += 1
            if count >= COMMIT_BUFFER_SIZE:
                count = 0
                self.kafka_consumer.commit()

    def _consume(self, msg, handler):
        if not self.cache:
            handler(msg)
            return
        key = f"{self.topic}:{config.GROUP_ID}:{msg.value.get('ID')}"
        if not self.cache.exists(key):
            handler(msg)
            self.cache.set(key, "")

    def close_consumer(self):
        try:
            self.kafka_consumer.close()
        except Exception:
            pass


def handle(msg):
    print(msg.value)


def consume(topic):
    consumer = Consumer(topic)
    try:
        consumer.consume(handle)
    except Exception as e:
        logger.exception(f"consume fail, err is {e}")
    finally:
        consumer.close_consumer()


def multi_consume(topic):
    partitions_num = 1
    kafka_consumer = KafkaConsumer(bootstrap_servers=config.SERVERS)
    if kafka_consumer.partitions_for_topic(topic):
        partitions_num = len(kafka_consumer.partitions_for_topic(topic))
    kafka_consumer.close()

    process_num = min(partitions_num, cpu_count() * 2)
    process = [
        Process(target=consume, args=(topic, )) for _ in range(process_num)
    ]
    for p in process:
        p.start()
    for p in process:
        p.join()


if __name__ == '__main__':
    topic = "test"
    consume(topic)
